'''
Author: Diana Tang
Date: 2025-04-04 22:40:02
LastEditors: Diana Tang
Description: some description
FilePath: /pub-scholar-assist/scripts/index_documents.py
'''
#!/usr/bin/env python
"""
文档索引脚本
用于将文档添加到知识库
"""

import os
import sys
import argparse
import codecs

# 添加项目根目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
sys.path.append(project_root)

# 设置默认编码为UTF-8
sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer)
sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer)

from app.model_loader import ModelLoader
from app.document_processor import DocumentProcessor
from app.vector_store import VectorStore

def index_documents(docs_dir, config_path="config/config.yaml"):
    """索引文档目录到向量存储"""
    print(f"开始索引文档目录: {docs_dir}")
    
    # 初始化ModelLoader并加载嵌入模型
    model_loader = ModelLoader(config_path)
    embedding_model = model_loader.load_embedding_model()
    
    # 初始化文档处理器
    doc_processor = DocumentProcessor(config_path)
    
    # 初始化向量存储
    vector_store = VectorStore(embedding_model, config_path)
    vector_store.load()  # 尝试加载现有索引
    
    # 处理文档
    chunks = doc_processor.process_directory(docs_dir)
    
    # 添加到向量存储
    vector_store.add_documents(chunks)
    
    # 保存索引
    vector_store.save()
    
    print(f"成功索引 {len(chunks)} 个文档块")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="索引文档到DeepSeek知识库")
    parser.add_argument("--docs", required=True, help="文档目录路径")
    parser.add_argument("--config", default="config/config.yaml", help="配置文件路径")
    
    args = parser.parse_args()
    
    index_documents(args.docs, args.config)